#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
🍪 YouTube Cookies 自动化设置脚本

此脚本提供多种方法自动获取和设置YouTube cookies，
方便进行YouTube视频下载测试。

使用方法:
    python setup_youtube_cookies.py
    
或指定特定方法:
    python setup_youtube_cookies.py --method chrome
    python setup_youtube_cookies.py --method firefox
    python setup_youtube_cookies.py --method auto
"""

import os
import sys
import json
import argparse
import subprocess
from pathlib import Path
from typing import Optional, Dict, List, Tuple

# 添加项目根目录到路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

try:
    import browser_cookie3
    BROWSER_COOKIE3_AVAILABLE = True
except ImportError:
    BROWSER_COOKIE3_AVAILABLE = False

try:
    import yt_dlp
    YT_DLP_AVAILABLE = True
except ImportError:
    YT_DLP_AVAILABLE = False


class YouTubeCookieSetup:
    """YouTube Cookies 自动化设置类"""
    
    def __init__(self):
        self.project_root = Path(__file__).parent
        self.cookies_file = self.project_root / "youtube_cookies.txt"
        self.test_url = "https://youtube.com/shorts/3598RsRvvIw?si=JeNy6GPe1sMeqXHD"
        
    def print_header(self):
        """打印脚本头部信息"""
        print("🍪" + "="*60)
        print("  YouTube Cookies 自动化设置脚本")
        print("="*62)
        print()
        
    def check_dependencies(self) -> Dict[str, bool]:
        """检查依赖项"""
        print("📋 检查依赖项...")
        
        deps = {
            'browser_cookie3': BROWSER_COOKIE3_AVAILABLE,
            'yt_dlp': YT_DLP_AVAILABLE,
        }
        
        for dep, available in deps.items():
            status = "✅" if available else "❌"
            print(f"  {status} {dep}: {'已安装' if available else '未安装'}")
            
        return deps
        
    def install_dependencies(self):
        """安装缺失的依赖"""
        deps = self.check_dependencies()
        missing = [dep for dep, available in deps.items() if not available]
        
        if not missing:
            print("✅ 所有依赖都已安装")
            return True
            
        print(f"\n📦 正在安装缺失的依赖: {', '.join(missing)}")
        
        try:
            for dep in missing:
                print(f"  安装 {dep}...")
                result = subprocess.run([
                    sys.executable, "-m", "pip", "install", dep
                ], capture_output=True, text=True)
                
                if result.returncode == 0:
                    print(f"  ✅ {dep} 安装成功")
                else:
                    print(f"  ❌ {dep} 安装失败: {result.stderr}")
                    return False
                    
            print("✅ 所有依赖安装完成")
            return True
            
        except Exception as e:
            print(f"❌ 安装依赖时发生错误: {e}")
            return False
            
    def get_cookies_from_browser(self, browser: str) -> Optional[List]:
        """从指定浏览器获取cookies"""
        if not BROWSER_COOKIE3_AVAILABLE:
            print("❌ browser_cookie3 未安装，无法自动获取cookies")
            return None
            
        print(f"🔍 正在从 {browser} 获取 YouTube cookies...")
        
        try:
            if browser.lower() == 'chrome':
                cj = browser_cookie3.chrome(domain_name='.youtube.com')
            elif browser.lower() == 'firefox':
                cj = browser_cookie3.firefox(domain_name='.youtube.com')
            elif browser.lower() == 'safari':
                cj = browser_cookie3.safari(domain_name='.youtube.com')
            elif browser.lower() == 'edge':
                cj = browser_cookie3.edge(domain_name='.youtube.com')
            else:
                print(f"❌ 不支持的浏览器: {browser}")
                return None
                
            cookies = list(cj)
            if cookies:
                print(f"✅ 从 {browser} 获取到 {len(cookies)} 个 cookies")
                return cookies
            else:
                print(f"⚠️  从 {browser} 未获取到 cookies")
                return None
                
        except Exception as e:
            print(f"❌ 从 {browser} 获取 cookies 失败: {e}")
            return None
            
    def save_cookies_to_file(self, cookies: List, filename: str = None) -> bool:
        """保存cookies到文件"""
        if filename is None:
            filename = self.cookies_file
            
        try:
            with open(filename, 'w', encoding='utf-8') as f:
                f.write("# Netscape HTTP Cookie File\n")
                f.write("# This is a generated file! Do not edit.\n\n")
                
                for cookie in cookies:
                    # 转换为Netscape格式
                    domain = cookie.domain
                    flag = "TRUE" if domain.startswith('.') else "FALSE"
                    path = cookie.path
                    secure = "TRUE" if cookie.secure else "FALSE"
                    expires = str(int(cookie.expires)) if cookie.expires else "0"
                    name = cookie.name
                    value = cookie.value
                    
                    f.write(f"{domain}\t{flag}\t{path}\t{secure}\t{expires}\t{name}\t{value}\n")
                    
            print(f"✅ Cookies 已保存到: {filename}")
            return True
            
        except Exception as e:
            print(f"❌ 保存 cookies 失败: {e}")
            return False
            
    def test_cookies_with_ytdlp(self, method: str = "file") -> bool:
        """使用yt-dlp测试cookies"""
        if not YT_DLP_AVAILABLE:
            print("❌ yt-dlp 未安装，无法测试cookies")
            return False
            
        print(f"🧪 使用 yt-dlp 测试 cookies ({method})...")
        
        try:
            if method == "file":
                cmd = [
                    "yt-dlp", 
                    "--cookies", str(self.cookies_file),
                    "--print-json",
                    "--quiet",
                    self.test_url
                ]
            elif method == "browser":
                cmd = [
                    "yt-dlp",
                    "--cookies-from-browser", "chrome",
                    "--print-json", 
                    "--quiet",
                    self.test_url
                ]
            else:
                print(f"❌ 不支持的测试方法: {method}")
                return False
                
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
            
            if result.returncode == 0:
                try:
                    video_info = json.loads(result.stdout)
                    title = video_info.get('title', 'Unknown')
                    duration = video_info.get('duration', 0)
                    print(f"✅ Cookies 测试成功!")
                    print(f"   视频标题: {title}")
                    print(f"   视频时长: {duration}秒")
                    return True
                except json.JSONDecodeError:
                    print("⚠️  获取到响应但解析失败")
                    return False
            else:
                error_msg = result.stderr.strip()
                if "Sign in to confirm" in error_msg:
                    print("❌ Cookies 无效，需要重新获取")
                else:
                    print(f"❌ 测试失败: {error_msg}")
                return False
                
        except subprocess.TimeoutExpired:
            print("❌ 测试超时")
            return False
        except Exception as e:
            print(f"❌ 测试时发生错误: {e}")
            return False
            
    def auto_setup(self) -> bool:
        """自动设置cookies"""
        print("🚀 开始自动设置 YouTube cookies...")
        
        # 尝试不同的浏览器
        browsers = ['chrome', 'firefox', 'safari', 'edge']
        
        for browser in browsers:
            print(f"\n📱 尝试从 {browser} 获取cookies...")
            cookies = self.get_cookies_from_browser(browser)
            
            if cookies:
                # 保存到文件
                if self.save_cookies_to_file(cookies):
                    # 测试cookies
                    if self.test_cookies_with_ytdlp("file"):
                        print(f"🎉 成功从 {browser} 设置 cookies!")
                        return True
                    else:
                        print(f"⚠️  从 {browser} 获取的 cookies 无效")
                        
        print("❌ 所有浏览器都无法获取有效的 cookies")
        return False
        
    def manual_guide(self):
        """显示手动设置指南"""
        print("\n📖 手动设置指南:")
        print("-" * 40)
        print("1. 浏览器扩展方法（推荐）:")
        print("   Chrome: 安装 'Get cookies.txt LOCALLY' 扩展")
        print("   Firefox: 安装 'cookies.txt' 扩展")
        print()
        print("2. yt-dlp 自动方法:")
        print("   yt-dlp --cookies-from-browser chrome --list-formats <URL>")
        print()
        print("3. 开发者工具方法:")
        print("   F12 → Application → Cookies → youtube.com")
        print()
        print(f"4. 将cookies文件保存到: {self.cookies_file}")
        print()
        
    def interactive_setup(self):
        """交互式设置"""
        print("\n🔧 交互式设置模式")
        print("-" * 30)
        
        while True:
            print("\n请选择操作:")
            print("1. 自动从浏览器获取cookies")
            print("2. 使用yt-dlp浏览器方法测试")
            print("3. 测试现有cookies文件")
            print("4. 显示手动设置指南")
            print("5. 退出")
            
            choice = input("\n请输入选择 (1-5): ").strip()
            
            if choice == '1':
                self.auto_setup()
            elif choice == '2':
                self.test_cookies_with_ytdlp("browser")
            elif choice == '3':
                if self.cookies_file.exists():
                    self.test_cookies_with_ytdlp("file")
                else:
                    print("❌ cookies文件不存在")
            elif choice == '4':
                self.manual_guide()
            elif choice == '5':
                print("👋 再见!")
                break
            else:
                print("❌ 无效选择，请重试")
                
    def run(self, method: str = "auto"):
        """运行脚本"""
        self.print_header()
        
        # 安装依赖
        if not self.install_dependencies():
            print("\n❌ 依赖安装失败，请手动安装后重试")
            return False
            
        print("\n" + "="*50)
        
        if method == "auto":
            success = self.auto_setup()
            if not success:
                print("\n自动设置失败，进入交互模式...")
                self.interactive_setup()
        elif method == "interactive":
            self.interactive_setup()
        elif method in ["chrome", "firefox", "safari", "edge"]:
            cookies = self.get_cookies_from_browser(method)
            if cookies:
                self.save_cookies_to_file(cookies)
                self.test_cookies_with_ytdlp("file")
        elif method == "test":
            if self.cookies_file.exists():
                self.test_cookies_with_ytdlp("file")
            else:
                print("❌ cookies文件不存在")
        elif method == "test-browser":
            self.test_cookies_with_ytdlp("browser")
        else:
            print(f"❌ 不支持的方法: {method}")
            
        return True


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description="YouTube Cookies 自动化设置脚本",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  python setup_youtube_cookies.py                    # 自动模式
  python setup_youtube_cookies.py --method chrome    # 从Chrome获取
  python setup_youtube_cookies.py --method firefox   # 从Firefox获取
  python setup_youtube_cookies.py --method test      # 测试现有cookies
  python setup_youtube_cookies.py --method interactive # 交互模式
        """
    )
    
    parser.add_argument(
        '--method', 
        choices=['auto', 'chrome', 'firefox', 'safari', 'edge', 'test', 'test-browser', 'interactive'],
        default='auto',
        help='设置方法 (默认: auto)'
    )
    
    args = parser.parse_args()
    
    try:
        setup = YouTubeCookieSetup()
        setup.run(args.method)
    except KeyboardInterrupt:
        print("\n\n👋 用户取消操作")
    except Exception as e:
        print(f"\n❌ 发生未知错误: {e}")
        return 1
        
    return 0


if __name__ == "__main__":
    sys.exit(main()) 